import torch
import torch.nn as nn
from ultralytics.nn.modules.conv import Conv


# 1111
class space_to_depth(nn.Module):
    # Changing the dimension of the Tensor
    def __init__(self, inc=64, dimension=1):
        super().__init__()
        self.d = dimension
        # self.conv = Conv(4*inc, 2*inc, 1, 1)

    def forward(self, x):
         # return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

         return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)

if __name__ == '__main__':
    input = torch.randn(1, 128, 8, 8)	# bachsize, c, h, w
    spd = space_to_depth()
    # spd = space_to_depth(128)
    output = spd(input)
    print(output.shape)	# torch.Size([1, 512, 4, 4])
